"""
快慢时间二维矩阵数据处理模块
负责从IQ数据中提取chirp数据并生成快慢时间二维矩阵
"""

import numpy as np
from typing import Dict, Any, Optional, Tuple
from src.logger import Logger


class FastSlowTimeProcessor:
    """快慢时间数据处理器"""
    
    def __init__(self, config: Dict[str, Any], logger: Optional[Logger] = None):
        """
        初始化快慢时间处理器
        
        Args:
            config: 快慢时间处理配置（需包含fmcw_signal配置）
            logger: 日志记录器
        """
        self.config = config
        self.logger = logger or Logger("fast_slow_time")
        
        # 从FMCW信号配置计算三角波时序参数
        fmcw_config = config.get('fmcw_signal', {})
        waveform_type = fmcw_config.get('waveform_type', 'linear')
        
        if waveform_type == 'triangular':
            # 三角波模式：根据fmcw_signal参数计算时序
            bandwidth = fmcw_config.get('bandwidth', 200e6)
            sweep_rate = fmcw_config.get('sweep_rate', 2e12)
            down_bandwidth = fmcw_config.get('down_bandwidth', 200e6)
            down_sweep_rate = fmcw_config.get('down_sweep_rate', 2e12)
            
            # 计算三角波时间参数
            self.up_sweep_duration = bandwidth / sweep_rate  # 上扫时间
            self.down_sweep_duration = down_bandwidth / down_sweep_rate  # 下扫时间
            self.triangle_period = self.up_sweep_duration + self.down_sweep_duration  # 总周期
            
            self.logger.info("三角波模式：从fmcw_signal配置计算时序参数")
            self.logger.info(f"  上扫时间: {self.up_sweep_duration*1e6:.1f} μs")
            self.logger.info(f"  下扫时间: {self.down_sweep_duration*1e6:.1f} μs")
            self.logger.info(f"  三角波周期: {self.triangle_period*1e6:.1f} μs")
        else:
            # 线性扫频模式：使用默认值
            self.triangle_period = 100e-6  # 默认100μs
            self.up_sweep_duration = 100e-6  # 默认100μs
            self.down_sweep_duration = 0  # 线性扫频无下降段
            
            self.logger.info("线性扫频模式：使用默认时序参数")
            self.logger.info(f"  扫频时间: {self.up_sweep_duration*1e6:.1f} μs")
        
        # 快时间窗口参数
        self.fast_time_start_offset = config.get('fast_time_start_offset', 5e-6)  # 5μs偏移
        self.fast_time_window = config.get('fast_time_window', 40e-6)  # 40μs窗口
        
        self.logger.info(f"快时间窗口: {self.fast_time_window*1e6:.1f} μs")
        self.logger.info(f"快时间起始偏移: {self.fast_time_start_offset*1e6:.1f} μs")
    
    def calculate_sampling_parameters(self, sample_rate: float) -> Dict[str, Any]:
        """
        计算采样参数
        
        Args:
            sample_rate: 采样率 (Hz)
            
        Returns:
            采样参数字典
        """
        # 计算每个样本的时间间隔
        sample_interval = 1.0 / sample_rate

        # 使用四舍五入将时间转换为样本点，避免系统性偏小（下取整）导致的累计误差
        triangle_period_samples = int(round(self.triangle_period / sample_interval))
        fast_time_window_samples = int(round(self.fast_time_window / sample_interval))
        fast_time_start_samples = int(round(self.fast_time_start_offset / sample_interval))

        # 基本健壮性保护：至少保证为1个样本
        triangle_period_samples = max(triangle_period_samples, 1)
        fast_time_window_samples = max(fast_time_window_samples, 1)

        # 如果起始偏移超过一个周期，进行折返并给出提示
        if fast_time_start_samples >= triangle_period_samples:
            self.logger.warning(
                "快时间起始样本数超过一个三角波周期，自动对周期取模处理"
            )
            fast_time_start_samples = fast_time_start_samples % triangle_period_samples

        # 如果窗口过长，确保窗口不会跨越本次上扫段（尽量不跨越一个周期）
        if fast_time_start_samples + fast_time_window_samples > triangle_period_samples:
            self.logger.warning(
                "快时间窗口超出单周期范围，已截断以避免跨周期"
            )
            fast_time_window_samples = max(1, triangle_period_samples - fast_time_start_samples)
        
        # 计算chirp步长（一个三角波周期的样本数）
        chirp_step = triangle_period_samples
        
        params = {
            'sample_interval': sample_interval,
            'triangle_period_samples': triangle_period_samples,
            'fast_time_window_samples': fast_time_window_samples,
            'fast_time_start_samples': fast_time_start_samples,
            'chirp_step': chirp_step,
            'effective_samples_per_chirp': fast_time_window_samples
        }
        
        self.logger.info("采样参数计算完成:")
        self.logger.info(f"  采样间隔: {sample_interval*1e6:.3f} μs")
        self.logger.info(f"  三角波周期样本数: {triangle_period_samples}")
        self.logger.info(f"  快时间窗口样本数: {fast_time_window_samples}")
        self.logger.info(f"  快时间起始样本数: {fast_time_start_samples}")
        self.logger.info(f"  Chirp步长: {chirp_step} 样本")
        
        return params
    
    def extract_chirp_matrix(self, iq_data: np.ndarray, sample_rate: float) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        从IQ数据中提取chirp二维矩阵
        
        Args:
            iq_data: 输入IQ数据
            sample_rate: 采样率
            
        Returns:
            (chirp_matrix, info) 元组
        """
        try:
            self.logger.info("开始提取chirp二维矩阵...")
            
            # 计算采样参数
            params = self.calculate_sampling_parameters(sample_rate)
            
            total_samples = len(iq_data)
            chirp_step = params['chirp_step']
            fast_time_window_samples = params['fast_time_window_samples']
            fast_time_start_samples = params['fast_time_start_samples']
            
            # 计算可提取的chirp数量
            max_chirp_start = total_samples - fast_time_start_samples - fast_time_window_samples
            num_chirps = max_chirp_start // chirp_step + 1
            
            if num_chirps <= 0:
                self.logger.error("数据长度不足以提取任何chirp")
                return np.array([]), {}
            
            # 创建chirp矩阵
            chirp_matrix = np.zeros((num_chirps, fast_time_window_samples), dtype=complex)
            
            # 提取每个chirp
            for chirp_idx in range(num_chirps):
                # 计算当前chirp的起始位置
                start_idx = fast_time_start_samples + chirp_idx * chirp_step
                end_idx = start_idx + fast_time_window_samples
                
                # 提取chirp数据
                chirp_matrix[chirp_idx, :] = iq_data[start_idx:end_idx]
            
            # 创建信息字典
            info = {
                'matrix_shape': chirp_matrix.shape,
                'num_chirps': num_chirps,
                'fast_time_samples': fast_time_window_samples,
                'total_input_samples': total_samples,
                # 未使用样本应以最后一个窗口的结束位置为基准，不应再多减去一个周期
                'unused_samples': int(
                    total_samples - (
                        fast_time_start_samples + (num_chirps - 1) * chirp_step + fast_time_window_samples
                    )
                ),
                'sampling_params': params,
                'extraction_config': {
                    'fast_time_start_offset_us': self.fast_time_start_offset * 1e6,
                    'fast_time_window_us': self.fast_time_window * 1e6,
                    'triangle_period_us': self.triangle_period * 1e6
                }
            }
            
            self.logger.info(f"成功提取chirp矩阵:")
            self.logger.info(f"  矩阵形状: {chirp_matrix.shape}")
            self.logger.info(f"  Chirp数量: {num_chirps}")
            self.logger.info(f"  快时间样本数: {fast_time_window_samples}")
            self.logger.info(f"  未使用样本数: {info['unused_samples']}")
            
            return chirp_matrix, info
            
        except Exception as e:
            self.logger.error(f"提取chirp矩阵失败: {e}")
            return np.array([]), {}
    
    def get_time_axes(self, chirp_matrix: np.ndarray, sample_rate: float, sampling_params: Optional[Dict[str, Any]] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        获取快慢时间轴
        
        Args:
            chirp_matrix: chirp矩阵
            sample_rate: 采样率
            
        Returns:
            (slow_time_axis, fast_time_axis) 元组
        """
        try:
            num_chirps, fast_time_samples = chirp_matrix.shape
            sample_interval = 1.0 / sample_rate

            # 快时间轴（微秒）
            fast_time_axis = np.arange(fast_time_samples) * sample_interval * 1e6

            # 慢时间轴（毫秒）应以实际提取步长计算：chirp_step * sample_interval
            if sampling_params is not None and 'chirp_step' in sampling_params:
                chirp_step = sampling_params['chirp_step']
                slow_step_s = chirp_step * sample_interval
            else:
                # 降级：使用配置的triangle_period（可能与实际步长略有出入）
                slow_step_s = self.triangle_period
            slow_time_axis = np.arange(num_chirps) * slow_step_s * 1e3
            
            self.logger.info(f"时间轴生成完成:")
            self.logger.info(f"  快时间轴: 0 - {fast_time_axis[-1]:.2f} μs, {len(fast_time_axis)} 点")
            self.logger.info(f"  慢时间轴: 0 - {slow_time_axis[-1]:.2f} ms, {len(slow_time_axis)} 点")
            
            return slow_time_axis, fast_time_axis
            
        except Exception as e:
            self.logger.error(f"生成时间轴失败: {e}")
            return np.array([]), np.array([])
    
    def validate_chirp_matrix(self, chirp_matrix: np.ndarray, info: Dict[str, Any]) -> bool:
        """
        验证chirp矩阵的有效性
        
        Args:
            chirp_matrix: chirp矩阵
            info: 提取信息
            
        Returns:
            验证是否通过
        """
        try:
            if chirp_matrix.size == 0:
                self.logger.error("chirp矩阵为空")
                return False
            
            if chirp_matrix.ndim != 2:
                self.logger.error(f"chirp矩阵维度错误: {chirp_matrix.ndim}")
                return False
            
            num_chirps, fast_time_samples = chirp_matrix.shape
            
            if num_chirps < 2:
                self.logger.warning(f"chirp数量较少: {num_chirps}")
            
            if fast_time_samples < 10:
                self.logger.warning(f"快时间样本数较少: {fast_time_samples}")
            
            # 检查是否有NaN或Inf值
            if np.any(np.isnan(chirp_matrix)):
                self.logger.error("chirp矩阵包含NaN值")
                return False
            
            if np.any(np.isinf(chirp_matrix)):
                self.logger.error("chirp矩阵包含Inf值")
                return False
            
            # 检查数据范围合理性
            magnitude = np.abs(chirp_matrix)
            if np.max(magnitude) > 10.0:
                self.logger.warning(f"chirp数据幅值较大: {np.max(magnitude):.3f}")
            
            self.logger.info("chirp矩阵验证通过")
            return True
            
        except Exception as e:
            self.logger.error(f"验证chirp矩阵失败: {e}")
            return False
    
    def process_iq_to_matrix(self, iq_data: np.ndarray, sample_rate: float) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        完整的IQ数据到chirp矩阵处理流程
        
        Args:
            iq_data: 输入IQ数据
            sample_rate: 采样率
            
        Returns:
            (chirp_matrix, processing_info) 元组
        """
        try:
            self.logger.info("=" * 60)
            self.logger.info("开始快慢时间二维矩阵处理")
            self.logger.info("=" * 60)
            
            # 提取chirp矩阵
            chirp_matrix, extract_info = self.extract_chirp_matrix(iq_data, sample_rate)
            
            if chirp_matrix.size == 0:
                self.logger.error("chirp矩阵提取失败")
                return np.array([]), {}
            
            # 验证chirp矩阵
            if not self.validate_chirp_matrix(chirp_matrix, extract_info):
                self.logger.error("chirp矩阵验证失败")
                return np.array([]), {}
            
            # 生成时间轴
            slow_time_axis, fast_time_axis = self.get_time_axes(
                chirp_matrix, sample_rate, sampling_params=extract_info.get('sampling_params')
            )
            
            # 创建完整处理信息
            processing_info = {
                'extraction_info': extract_info,
                'time_axes': {
                    'slow_time_axis_ms': slow_time_axis,
                    'fast_time_axis_us': fast_time_axis
                },
                'data_statistics': {
                    'matrix_shape': chirp_matrix.shape,
                    'total_samples': chirp_matrix.size,
                    'data_type': str(chirp_matrix.dtype),
                    'magnitude_stats': {
                        'min': float(np.min(np.abs(chirp_matrix))),
                        'max': float(np.max(np.abs(chirp_matrix))),
                        'mean': float(np.mean(np.abs(chirp_matrix))),
                        'std': float(np.std(np.abs(chirp_matrix)))
                    }
                },
                'processing_config': self.config
            }
            
            self.logger.info("=" * 60)
            self.logger.info("快慢时间二维矩阵处理完成")
            self.logger.info("=" * 60)
            self.logger.info(f"最终矩阵形状: {chirp_matrix.shape}")
            self.logger.info(f"数据类型: {chirp_matrix.dtype}")
            self.logger.info(f"总样本数: {chirp_matrix.size}")
            
            return chirp_matrix, processing_info
            
        except Exception as e:
            self.logger.error(f"处理IQ数据到矩阵失败: {e}")
            return np.array([]), {}
